from langdev_helper.llm.qwen import llm

from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory(return_messages=True)
memory.load_memory_variables({})

memory.save_context({"input": "我的名字叫张三"}, {"output": "你好，张三"})
memory.load_memory_variables({})

memory.save_context({"input": "我是一名 IT 程序员"}, {"output": "好的，我知道了"})
memory.load_memory_variables({})

from langchain.prompts import ChatPromptTemplate
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个乐于助人的助手。"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{user_input}"),
    ]
)
chain = prompt | llm

user_input = "你知道我的名字吗？"
print('Q:', user_input)
history = memory.load_memory_variables({})["history"]

res = chain.invoke({"user_input": user_input, "history": history})
print(res.content)

user_input = "中国最高的山是什么山？"
print('Q:', user_input)
res = chain.invoke({"user_input": user_input, "history": history})
print(res.content)
memory.save_context({"input": user_input}, {"output": res.content})


res = chain.invoke({"user_input": "我们聊得最后一个问题是什么？", "history": history})
print(res.content)
